#ifndef __BWLIB_NUMBER_THEORY_SIEVE_H__
#define __BWLIB_NUMBER_THEORY_SIEVE_H__
#include "../config.h"
#include <iterator>

namespace bwlib
{
	/* @brief: 泛型线性筛法
	   @complexity: O(n)
	   @memory: O(N + 1)
	   @param: N, 需要筛的最大数 
	   		   [out] prime, 质数数组
			   prime_end, 质数数组的末尾，
			   		如果质数超过 prime_end - prime，算法停止
			   single_prime, 仿函数
			   		@brief: 遇到一个素数时调用
					@param: p，素数
			   first_factor, 仿函数
			   		@brief: 当 !(p | i) 时调用，p 是素数
					@param: i * p, i, p
			   more_factor, 仿函数
			   		@brief: 当 p | i 时调用，p 是素数
					@param: i * p, i, p
			   not_prime, 临时空间，表示 i 是否是素数
	   @return: 素数数组最后一个元素的下一个位置 */
	template<typename RandomIter,
		typename SinglePrime, 
		typename FirstFactor, 
		typename MoreFactor>
	inline RandomIter linear_sieves(
		typename std::iterator_traits<
			RandomIter>::value_type N, 
		RandomIter prime, 
		RandomIter prime_end, 
		SinglePrime single_prime, 
		FirstFactor first_factor, 
		MoreFactor more_factor, 
		bool* temp = 0)
	{
		typedef typename std::iterator_traits<
			RandomIter>::value_type value_type;

		bool new_memory = false;
		if(temp == 0)
		{
			new_memory = true;
			temp = new bool[N + 1];
		}

		for(int i = 0; i <= N; ++i)
			temp[i] = true;

		temp[0] = false;
		temp[1] = false;
		int total = 0;
		for(value_type i = 2; i <= N; ++i)
		{
			if(temp[i] == true)
			{
				prime[total++] = i;
				single_prime(i);
				if(prime + total == prime_end)
					break;
			}

			for(int j = 0; j != total; ++j)
			{
				value_type t = prime[j] * i;
				if(t > N) break;
				temp[t] = false;
				if(i % prime[j] == 0)
				{
					more_factor(t, i, prime[j]);
					break;
				} else {
					first_factor(t, i, prime[j]);
				}
			}
		}

		if(new_memory)
			delete[] temp;

		return prime + total;
	}

	namespace detail
	{
		struct null_factor
		{
			template<typename T>
			void operator() (T) {}

			template<typename T1, typename T2, typename T3>
			void operator() (T1, T2, T3) {}
		};

		// functor for phi(x)
		template<typename Iter>
		class phi_prime
		{
			Iter phi;
		public:
			phi_prime(Iter phi_iter)
				: phi(phi_iter) {}

			template<typename T>
			void operator() (const T& p)
			{
				phi[p] = p - 1;
			}
		};

		template<typename Iter>
		class phi_first_factor
		{
			Iter phi;
		public:
			phi_first_factor(Iter phi_iter)
				: phi(phi_iter) {}

			template<typename T>
			void operator() (const T& t, const T& d, const T& p)
			{
				phi[t] = phi[d] * (p - 1);
			}
		};

		template<typename Iter>
		class phi_more_factor
		{
			Iter phi;
		public:
			phi_more_factor(Iter phi_iter)
				: phi(phi_iter) {}

			template<typename T>
			void operator() (const T& t, const T& d, const T& p)
			{
				phi[t] = phi[d] * p;
			}
		};
		
		// functor for mobius function (miu)
		template<typename Iter>
		class mobius_prime
		{
			Iter mobius;
		public:
			mobius_prime(Iter mobius_iter)
				: mobius(mobius_iter) {}

			template<typename T>
			void operator() (const T& p)
			{
				mobius[p] = -1;
			}
		};

		template<typename Iter>
		class mobius_first_factor
		{
			Iter mobius;
		public:
			mobius_first_factor(Iter mobius_iter)
				: mobius(mobius_iter) {}

			template<typename T>
			void operator() (const T& t, const T& d, const T&)
			{
				mobius[t] = -mobius[d];
			}
		};

		template<typename Iter>
		class mobius_more_factor
		{
			Iter mobius;
		public:
			mobius_more_factor(Iter mobius_iter)
				: mobius(mobius_iter) {}

			template<typename T>
			void operator() (const T& t, const T&, const T&)
			{
				mobius[t] = 0;
			}
		};
	}

	/* @prime: 线性筛计算素数
	   @complexity: O(n)
	   @memory: O(N + 1)
	   @param: N, 需要筛的最大数 
	   		   [out] prime, 素数数组
			   prime_end, 素数数组末尾
			   not_prime, 临时空间
	   @return: 素数数组最后一个元素的下一个位置 */
	template<typename RandomIter>
	inline RandomIter linear_prime_sieves(
		typename std::iterator_traits<
			RandomIter>::value_type N, 
		RandomIter prime,
		RandomIter prime_end, 
		bool* temp = 0)
	{
		detail::null_factor nf;
		return linear_sieves(N, prime, prime_end, nf, nf, nf, temp);
	}

	/* @prime: 线性筛计算欧拉函数
	   @complexity: O(n)
	   @memory: O(N + 1)
	   @param: N, 需要筛的最大数 
	   		   [out] prime, 素数数组
			   prime_end, 素数数组末尾
			   [out] phi, 欧拉函数值
			   not_prime, 临时空间
	   @return: 素数数组最后一个元素的下一个位置 */
	template<typename RandomIter, typename RandomPhiIter>
	inline RandomIter linear_phi_sieves(
		typename std::iterator_traits<
			RandomIter>::value_type N, 
		RandomIter prime, 
		RandomIter prime_end, 
		RandomPhiIter phi,
		bool* temp = 0)
	{
		phi[1] = 1;
		return linear_sieves(N, prime, prime_end,
			detail::phi_prime<RandomIter>(phi), 
			detail::phi_first_factor<RandomIter>(phi), 
			detail::phi_more_factor<RandomIter>(phi), temp);
	}

	/* @prime: 线性筛计算莫比乌斯函数
	   @complexity: O(n)
	   @memory: O(N + 1)
	   @param: N, 需要筛的最大数 
	   		   [out] prime, 素数数组
			   prime_end, 素数数组末尾
			   [out] mobius, 莫比乌斯函数值
			   not_prime, 临时空间
	   @return: 素数数组最后一个元素的下一个位置 */
	template<typename RandomIter, typename RandomMobiusIter>
	inline RandomIter linear_mobius_sieves(
		typename std::iterator_traits<
			RandomIter>::value_type N, 
		RandomIter prime, 
		RandomIter prime_end, 
		RandomMobiusIter mobius,
		bool* temp = 0)
	{
		mobius[1] = 1;
		return linear_sieves(N, prime, prime_end,
			detail::mobius_prime<RandomIter>(mobius), 
			detail::mobius_first_factor<RandomIter>(mobius), 
			detail::mobius_more_factor<RandomIter>(mobius), temp);
	}
}

#endif
